from sklearn.ensemble import ExtraTreesRegressor

from my_models.base import Base



class TreeRegressor(Base):

    def __init__(self):

        # import ipdb;ipdb.set_trace();
        self.regressor = ExtraTreesRegressor(
                                                n_estimators=100,       # 决策树的数量
                                                max_depth=None,         # 树的最大深度，默认为None表示不限制
                                                min_samples_split=2,    # 分裂内部节点所需的最小样本数
                                                min_samples_leaf=1,     # 叶节点所需的最小样本数
                                                max_features='sqrt',    # 考虑用于分裂的最佳特征数量
                                                bootstrap=False,        # 是否使用有放回抽样
                                                n_jobs=-1,              # 并行作业数量，-1 表示使用所有可用核心
                                                random_state=42         # 随机种子，保证结果可复现
                                            )

    def train(self, x_train, y_train):
        self.clf = self.regressor.fit(x_train, y_train)

    def valid(self, x_test, y_test):
        score_c = self.clf.score(x_test, y_test)
        print( "TreeRegressor:{}".format(score_c))

    def save_model(self, save_path):
        pass

    def load_model(self, model_path):
        pass
